#include <memory>
#include <chrono>
#include <cstring>

#include "core/pcap_device.h"
#include "utils/logger.h"

PcapDevice::PcapDevice() : m_handle(nullptr), m_isActivated(false), m_isCapturing(false) {}

PcapDevice::~PcapDevice() { close(); }

bool PcapDevice::open(const std::string& deviceName) {
    if (m_handle) {
        LOG_WARNING("Device is already opened: " << m_deviceName);
        return true;
    }

    char errbuf[PCAP_ERRBUF_SIZE];
    /// Use newer version pcap_create, which it can set the device laterly.
    /// After pcap_activate, it is really opened.
    m_handle = pcap_create(deviceName.c_str(), errbuf);

    if (!m_handle) {
        LOG_ERROR("Failed to create device: " << deviceName);
        return false;
    }

    m_deviceName = deviceName;
    LOG_INFO("Successfully create device: " << deviceName);
    return true;
}

void PcapDevice::close() {
    if (m_handle) {
        stopCapture();
        destroy();
    }
    m_deviceName.clear();
}

bool PcapDevice::isOpen() const { return m_handle != nullptr; }

bool PcapDevice::startCapture() {
    std::lock_guard<std::mutex> lock(m_captureMutex);

    if (!m_handle) {
        LOG_ERROR("Cannot start capture - device not created.");;
        return false;
    }

    if (m_isCapturing.load()) {
        LOG_WARNING("Capture already started on device: " << m_deviceName);
        return true;
    } else {
        m_isCapturing.store(true);
    }

    if (!m_callback) {
        LOG_ERROR("Callback function of device " << m_deviceName << " not set");
        return false;
    }

    if (!m_isActivated) {
        activate();
    }

    LOG_DEBUG("Start capture packets on: " << m_deviceName);
    /// start capture thread
    try {
        m_captureThread = std::thread(&PcapDevice::captureLoop, this);
    } catch (const std::exception& e) {
        m_isCapturing.store(false);
        LOG_ERROR("Thread creation failed " << e.what() << " device: " << m_deviceName);
        return false;
    }

    return true;
}

void PcapDevice::stopCapture() {
    /// First set the stop flag to avoid deadlock.
    m_isCapturing.store(false);
    if (m_handle) {
        pcap_breakloop(m_handle);
    }

    std::lock_guard<std::mutex> lock(m_captureMutex);

    if (m_captureThread.joinable()) {
        try {
            m_captureThread.join();
        } catch (const std::exception& e) {
            LOG_ERROR("Thread join failed: " << e.what() << " device: " << m_deviceName);
        }
    }

    /// Reset statistics.
    m_packetsReceived.store(0);
    m_bytesReceived.store(0);
    m_packetsDropped.store(0);
}

bool PcapDevice::injectPacket(const uint8_t* data, size_t length) {
    if (!m_handle) {
        LOG_ERROR("Cannot inject packet - device not created.");
        return false;
    }

    if (pcap_inject(m_handle, data, length) == -1) {
        LOG_ERROR("Failed to inject packet: " << pcap_geterr(m_handle));
        return false;
    }

    return true;
}

bool PcapDevice::injectPacket(const Packet& packet) {
    return injectPacket(packet.data(), packet.length());
}

bool PcapDevice::setBPFFilter(const std::string& filterExpression) {
    if (!m_handle) {
        LOG_ERROR("Cannot set BPF filter - device not opened");
    }

    /// Save filter expression.
    m_BPFFilter = filterExpression;

    /// If the device is activated, set the filter immediately.
    if (m_isActivated) {
        return applyBPFFilter();
    }

    LOG_DEBUG("BPF filter will be set after device activation: " << filterExpression);
    return true;
}

bool PcapDevice::applyBPFFilter() {
    if (m_BPFFilter.empty()) {
        return true;
    }

    struct bpf_program fp;

    if (pcap_compile(m_handle, &fp, m_BPFFilter.c_str(), 0, PCAP_NETMASK_UNKNOWN) == -1) {
        LOG_ERROR("Cannot parse filter '" << m_BPFFilter << "': " << pcap_geterr(m_handle));
        return false;
    }

    if (pcap_setfilter(m_handle, &fp) == -1) {
        LOG_ERROR("Cannot install filter '" << m_BPFFilter << "': " << pcap_geterr(m_handle));
        return false;
    }

    pcap_freecode(&fp);

    LOG_INFO("Set BPF filter: " << m_BPFFilter << " on device: " << m_deviceName);
    return true;
}

bool PcapDevice::setPromiscuousMode(bool enabled) {
    if (!m_handle) {
        LOG_ERROR("Cannot set promiscuous mode - device not created.");
        return false;
    }

    if (pcap_set_promisc(m_handle, enabled ? 1 : 0) != 0) {
        LOG_ERROR("Failed to set promiscuous mode: " << pcap_geterr(m_handle));
        return false;
    }

    LOG_INFO("Promiscuous mode " << (enabled ? "enabled" : "disabled") << " on: " << m_deviceName);
    return true;
}

bool PcapDevice::setTimeout(int milliseconds) {
    if (!m_handle) {
        LOG_ERROR("Cannot set timeout - device not created.");
        return false;
    }

    if (pcap_set_timeout(m_handle, milliseconds) != 0) {
        LOG_ERROR("Failed to set timeout: " << pcap_geterr(m_handle));
        return false;
    }

    return true;
}

bool PcapDevice::setBufferSize(int size) {
    if (!m_handle) {
        LOG_ERROR("Cannot set buffer size - device not created.");
        return false;
    }

    if (pcap_set_buffer_size(m_handle, size) != 0) {
        LOG_ERROR("Failed to set buffer size: " << pcap_geterr(m_handle));
        return false;
    }

    return true;
}

uint64_t PcapDevice::getPacketReceived() const {
    return m_packetsReceived.load();
}

uint64_t PcapDevice::getPacketDropped() const {
    return m_packetsDropped.load();
}

uint64_t PcapDevice::getBytesReceived() const {
    return m_bytesReceived.load();
}

std::string PcapDevice::getName() const { return m_deviceName; }

std::string PcapDevice::getDescription() const {
    if (!m_handle) {
        return "Device not opened.";
    }

    pcap_if_t* alldevs;
    char errbuf[PCAP_ERRBUF_SIZE];

    if (pcap_findalldevs(&alldevs, errbuf) == -1) {
        return "Unknown device.";
    }

    for (pcap_if_t* d = alldevs; d; d = d->next) {
        if (d->name && m_deviceName == d->name) {
            std::string description = d->description ? d->description : "No description.";
            pcap_freealldevs(alldevs);
            return description;
        }
    }

    pcap_freealldevs(alldevs);
    return "Device not found.";
}

bool PcapDevice::setFilter(const std::string& filterExpression) {
    if (!m_handle) {
        LOG_ERROR("Cannot set filter - device not created.");
        return false;
    }

    struct bpf_program fp;
    if (pcap_compile(m_handle, &fp, filterExpression.c_str(), 0, PCAP_NETMASK_UNKNOWN) == -1) {
        LOG_ERROR("Couldn't parse filter " << filterExpression << " : " << pcap_geterr(m_handle));
        return false;
    }

    if (pcap_setfilter(m_handle, &fp) == -1) {
        LOG_ERROR("Couldn't install filter " << filterExpression << " : " << pcap_geterr(m_handle));
        return false;
    }

    pcap_freecode(&fp);
    LOG_INFO("Set BPF filter: " << filterExpression << " on: " << m_deviceName);
    return true;
}

void PcapDevice::setPacketCallback(PacketCallBack callback) {
    m_callback = std::move(callback);
}

bool PcapDevice::setSnapshotLength(int snapshot_len) {
    if (!m_handle) {
        LOG_ERROR("Cannot set snaplen - device not created.");
        return false;
    }

    if (pcap_set_snaplen(m_handle, snapshot_len) != 0) {
        LOG_ERROR("Failed to set snaplen: " << pcap_geterr(m_handle));
        return false;
    }
    return true;
}

bool PcapDevice::setImmediateMode(bool enabled) {
    if (!m_handle) {
        LOG_ERROR("Cannot set immediate mode - device not created.");
        return false;
    }

    if (pcap_set_immediate_mode(m_handle, enabled) != 0) {
        LOG_ERROR("Failed to set immediate mode : " << pcap_geterr(m_handle));
        return false;
    }
    return true;
}

bool PcapDevice::setRFMonitorMode(bool enabled) {
    if (!m_handle) {
        LOG_ERROR("Cannot set wireless monitor mode - device not created.");
        return false;
    }

    if (pcap_can_set_rfmon(m_handle) == 0){
        if (pcap_set_rfmon(m_handle, enabled) != 0) {
            LOG_ERROR("Failed to set wireless monitor mode: " << pcap_geterr(m_handle));
            return false;
        }
    } else {
        LOG_ERROR("Cannot set wireless monitor mode device not support.");
        return false;
    }
    return true;
}

bool PcapDevice::activate() {
    if (!m_handle) {
        LOG_ERROR("Cannot activate device " << m_deviceName << " - device not created.");
        return false;
    }

    if(pcap_activate(m_handle) != 0) {
        LOG_ERROR("Device " << m_deviceName << " failed to activate: " << pcap_geterr(m_handle));
        return false;
    }
    m_isActivated = true;
    LOG_DEBUG("Device " << m_deviceName << " activated");

    applyBPFFilter();
    return true;
}

void PcapDevice::captureLoop() {
    // For statistics.
    uint64_t batchPackets = 0;
    uint64_t batchBytes = 0;
    uint64_t batchDropped = 0;

    while (m_isCapturing.load()) {
        struct pcap_pkthdr* header;
        const u_char* packetData;

        int result = pcap_next_ex(m_handle, &header, &packetData);

        if (result == 1) {
            /// Successfully captured packet.
            if (packetData && m_callback) {
                Packet packet(packetData, header->len);
                if (packet.parse()) {
                    /// The callback processes exception.
                    try {
                        m_callback(std::move(packet));
                    } catch (const std::exception& e) {
                        LOG_ERROR("Callback error: " << e.what());
                    } catch (...) {
                        LOG_ERROR("Unknown callback error.");
                    }

                    ++batchPackets;
                    batchBytes += header->len;
                } else {
                    ++batchDropped;
                }
            }
            /// Batch update statistics.
            if (batchPackets >= 100) {
                m_packetsReceived.fetch_add(batchPackets);
                m_bytesReceived.fetch_add(batchBytes);
                m_packetsDropped.fetch_add(batchDropped);
                batchPackets = batchBytes = batchDropped = 0;
            }
        } else if (result == 0) {
            /// Timeout
            continue;
        } else if (result == -1) {
            /// Error
            LOG_ERROR("Capture error: " << pcap_geterr(m_handle));
            break;
        } else if (result == -2) {
            // End of off line file. Currently not support.
            break;
        }
    }

    if (batchPackets > 0) {
        m_packetsReceived.fetch_add(batchPackets);
        m_bytesReceived.fetch_add(batchBytes);
        m_packetsDropped.fetch_add(batchDropped);
    }
}

bool PcapDevice::cleanup() {
    if (m_handle) {
        pcap_close(m_handle);
        m_handle = nullptr;
    }
    m_isActivated = false;
    return true;
}

bool PcapDevice::destroy() {
    if (m_isCapturing) {
        stopCapture();
    }
    cleanup();
    LOG_DEBUG("Destroyed device: " << m_deviceName);
    return true;
}

void PcapDevice::packetHandler(u_char* userData, const struct pcap_pkthdr* header, const u_char* packetData) {
    /// Use for pcap_loop, but we use a different way in startCapture.
}
